# code: utf-8
import copy
import itertools as it

import numpy as np
from numpy import linalg as LA

import detail_design.models.detail_design.jxq.group as group
import detail_design.models.detail_design.jxq.model_operator as model_operator
from detail_design.models.detail_design.jxq.Sympgroup import SymBody


# 自实现梯度下降法

# 调整RPS点的顺序，与距离最近的clamp点靠近
def modify_rps_order(obj,rps_coor):
    if len(rps_coor)>1:
        clamp_points=get_clamp_point(obj)
        clamp_points_list=[v[0] for v in clamp_points]
        rps_combination_dict=[]
        rps_dist_list=[]
        for temp_rps in it.permutations(rps_coor,len(rps_coor)):
            dis_rps_clamp=0.0
            for clamp_point,rps_point in zip(clamp_points_list,temp_rps):
                dis_rps_clamp+=LA.norm(clamp_point-rps_point)
            rps_dist_list.append(dis_rps_clamp)
            rps_combination_dict.append(temp_rps)
        min_index = rps_dist_list.index(min(rps_dist_list))
        rps_coor=rps_combination_dict[min_index]
    print("调整后的rps坐标为：", rps_coor)
    return rps_coor


# 获取夹具点的坐标
def get_clamp_point(clamp=None):
    point = []
    clamp_point=clamp.obj.get('clamp_points', [])
    for i in range(len(clamp_point)):
        for j in range(len(clamp_point[i])):

            point.append(clamp.get_point(clamp_point[i][j]))
    # print(point)
    points = np.array(point).reshape(len(clamp_point), len(clamp_point[0]), 3)
    # print("夹紧点的坐标为：", points)
    return points




def get_a_new_example(obj):
    temp_obj=copy.deepcopy(obj)
    para_dict=copy.deepcopy(obj.obj['args'])
    new_para_dict=copy.deepcopy(para_dict)
    for key in para_dict:
        new_para_dict[key]['v']=key['v']+random.random()
    temp_obj.obj['args']=new_para_dict
    new_obj = group.Body(temp_obj.obj)
    new_obj = model_operator.update_model_by_history(new_obj, obj)
    return new_obj

def get_a_batch_of_examples(obj,n):
    result=[]
    for i in range(n):
        new_obj=get_a_new_example(obj)
        result.append(new_obj)
    return result


# 获得各个RPS点与夹具点之间的距离总和
def get_rps_diff(possible_para_sym_list,possible_para_dict, express):
    possible_value=[possible_para_dict[e[1]]['v'] for e in possible_para_sym_list]
    total_loss_sym=[e(possible_value) for e in express]
    total_loss=[LA.norm(np.array(e,dtype=float)) for e in total_loss_sym]
    return sum(total_loss)


# 夹具各部件尺寸更新
def update_args(obj, rps_coor, lr=0.05, dx=0.1, iters=40):
    import sympy as sym

    grad_dict = {}
    sym_obj=SymBody(obj.obj)

    sym_obj = model_operator.update_model_by_history(sym_obj, obj)
    possible_para_sym_list = [(sym_obj.syms[e],e) for e in sym_obj.args]
    possible_para_sym=[e[0] for e in possible_para_sym_list]
    possible_para_dict=copy.deepcopy(sym_obj.args)
    possible_para_list=copy.deepcopy(sym_obj.obj['args'])
    clamp_points = get_clamp_point(sym_obj)

    test_loss_sym = [sym.sqrt(sum((clamp_point - rps_point) * (clamp_point - rps_point).T)) for
                     clamp_point_list, rps_point in zip(clamp_points, rps_coor) for clamp_point in clamp_point_list]

    total_loss_sym=sum(test_loss_sym)
    total_loss_fun=sym.lambdify(([possible_para_sym]), total_loss_sym, 'numpy')


    for i in range(iters):
        possible_value = [possible_para_dict[e[1]]['v'] for e in possible_para_sym_list]
        base_loss=total_loss_fun(possible_value)
        # base_loss=get_rps_diff(possible_para_sym_list,possible_para_dict,total_loss_sym)

        if base_loss<1e-5:
            break
        for ele in possible_para_list:
            if 'isconst' not in ele or not ele['isconst']:
                possible_para_dict[ele['name']]['v'] += dx
                possible_value = [possible_para_dict[e[1]]['v'] for e in possible_para_sym_list]
                after_loss=total_loss_fun(possible_value)
                # after_loss = get_rps_diff(possible_para_sym_list,possible_para_dict,total_loss_sym)

                grad_temp = (after_loss - base_loss) / dx
                if abs(grad_temp) < 1e-6:
                    grad_dict[ele['name']] = 0
                else:
                    if abs(base_loss / grad_temp) < 5:
                        grad_dict[ele['name']] = base_loss / grad_temp
                    else:
                        grad_dict[ele['name']] = grad_temp
                possible_para_dict[ele['name']]['v'] -= dx

        for ele in possible_para_list:
            if ele['name'] in grad_dict:
                possible_para_dict[ele['name']]['v'] = possible_para_dict[ele['name']]['v'] - lr * grad_dict[ele['name']]
                if 'range' in ele:
                    if possible_para_dict[ele['name']]['v'] < ele['range'][0]:
                        possible_para_dict[ele['name']]['v'] = ele['range'][0]+0.5
                    if possible_para_dict[ele['name']]['v']> ele['range'][1]:
                        possible_para_dict[ele['name']]['v'] = ele['range'][1]

    new_para=[{'name':ele['name'],'v':possible_para_dict[ele['name']]['v']} for ele in possible_para_list]

    
    obj.obj['args']=new_para
    new_obj=group.Body(obj.obj)
    obj=model_operator.update_model_by_history(new_obj,obj)
    return obj

def get_rps_diff(obj, rps_coor):
    total_loss = 0
    clamp_points=get_clamp_point(obj)
    for clamp_point_list,rps_point in zip(clamp_points,rps_coor):
        new_rps_list=[rps_point for i in range(len(clamp_point_list))]
        for clamp_point, rps_point in zip(clamp_point_list, new_rps_list):
            total_loss += LA.norm(clamp_point - rps_point)
    return total_loss


    
# 模型更新
def model_update(obj, rps_coor, env=None, section=None, air_point=None, lr=1, dx=1, iters=1):

    # 调整RPS点的顺序，与距离最近的clamp点靠近
    rps_coor=modify_rps_order(obj,rps_coor)
    # rps_coor = rps_modify(obj, rps_coor)
    print("调整后的rps_coor：", rps_coor)

    obj = update_args(obj, rps_coor, lr=lr, dx=dx, iters=iters)
    print("更新后夹具点的坐标为：{}，距离为：{}".format(get_clamp_point(obj), get_rps_diff(obj, rps_coor)))
    print("更新后的零件尺寸为：", obj.args)
    # print("模型更新时间为：", time.time() - start_time)
    return obj




if __name__ == "__main__":
    rps1 = [[-370.5, 0, -95]]
    #
    # rps2 = [
    #         [-370.5, 80, -95],
    #         [-370.5, -80, -95],
    #     ]
    #
    # rps3 = [
    #
    #         [-185.0,  0, -90],
    #         [-360.5, 80, -95],
    #         [-360.5, -80, -95],
    #     ]
    # # Env = trimesh.load("work_2_xy.stl")
    # x = group.Group("clamp")
    # # x.show()
    # section = [1, 0, 0]
    # air_point = [0, 0, 0]
    # model_update(obj=x, rps_coor=rps1, env=None, section=None, air_point=None, lr=1, dx=0.5, iters=100)




